{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import necessary libraries\n",
    "\n",
    "import random\n",
    "import os, pickle\n",
    "import numpy as np\n",
    "import SimpleITK as sitk\n",
    "from collections import OrderedDict\n",
    "from scipy.ndimage import binary_fill_holes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# necessary utility functions\n",
    "\n",
    "def get_bbox(inp):\n",
    "    coords = np.where(inp != 0)\n",
    "    minz = np.min(coords[0])\n",
    "    maxz = np.max(coords[0]) + 1\n",
    "    minx = np.min(coords[1])\n",
    "    maxx = np.max(coords[1]) + 1\n",
    "    miny = np.min(coords[2])\n",
    "    maxy = np.max(coords[2]) + 1\n",
    "    return slice(minz, maxz), slice(minx, maxx), slice(miny, maxy)\n",
    "\n",
    "def convert_seg(seg):\n",
    "    \"\"\" convert brats labels from {0, 1, 2, 4} to {0, 1, 2, 3} \"\"\"\n",
    "    new_seg = np.zeros_like(seg)\n",
    "    new_seg[seg == 4] = 3\n",
    "    new_seg[seg == 2] = 1\n",
    "    new_seg[seg == 1] = 2\n",
    "    return new_seg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_path = '../data/raw'\n",
    "out_path = 'data/processed'\n",
    "\n",
    "names = os.listdir(data_path)\n",
    "for name in names:\n",
    "    flair = sitk.GetArrayFromImage(sitk.ReadImage(os.path.join(data_path, name, f'{name}_flair.nii.gz')))\n",
    "    t1 = sitk.GetArrayFromImage(sitk.ReadImage(os.path.join(data_path, name, f'{name}_t1.nii.gz')))\n",
    "    t1ce = sitk.GetArrayFromImage(sitk.ReadImage(os.path.join(data_path, name, f'{name}_t1ce.nii.gz')))\n",
    "    t2 = sitk.GetArrayFromImage(sitk.ReadImage(os.path.join(data_path, name, f'{name}_t2.nii.gz')))\n",
    "    seg = sitk.GetArrayFromImage(sitk.ReadImage(os.path.join(data_path, name, f'{name}_seg.nii.gz')))\n",
    "    img = np.stack([flair, t1, t1ce, t2]).astype(np.float32)\n",
    "    seg = convert_seg(seg)\n",
    "    # crop foreground regions\n",
    "    mask = np.zeros_like(seg).astype(bool)\n",
    "    for i in range(len(img)):\n",
    "        mask = mask | (img[i] != 0)\n",
    "    mask = binary_fill_holes(mask)\n",
    "    bbox = get_bbox(mask)\n",
    "    img = img[:, bbox[0], bbox[1], bbox[2]]\n",
    "    seg = seg[bbox[0], bbox[1], bbox[2]]\n",
    "    mask = mask[bbox[0], bbox[1], bbox[2]]\n",
    "    # normalization\n",
    "    for i in range(len(img)):\n",
    "        img[i][mask] = (img[i][mask] - img[i][mask].min()) / (img[i][mask].max() - img[i][mask].min())\n",
    "        img[i][mask == 0] = 0\n",
    "    # compensate label imbalance\n",
    "    approx_nsamp = 10000\n",
    "    samp_locs = OrderedDict()\n",
    "    for cls in [1, 2, 3]:\n",
    "        locs = np.argwhere(seg == cls)\n",
    "        nsamp = min(approx_nsamp, len(locs))\n",
    "        nsamp = max(nsamp, int(np.ceil(0.1 * len(locs))))\n",
    "        samp = locs[random.sample(range(len(locs)), nsamp)]\n",
    "        if len(samp) != 0:\n",
    "            samp_locs[cls] = samp\n",
    "    data = np.concatenate([img, seg[None]])\n",
    "    np.save(os.path.join(out_path, f'{name}.npy'), data)\n",
    "    with open(os.path.join(out_path, f'{name}.pkl'), 'wb') as f:\n",
    "        pickle.dump(samp_locs, f)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
